package saramsx

import (
	"context"
	"encoding/json"
	"gitee.com/chicken-c/go-learn/webook/pkg/logger"
	"github.com/IBM/sarama"
	"time"
)

type BatchHandler[T any] struct {
	l  logger.Logger
	fn func(mes []*sarama.ConsumerMessage, event []T) error
}

func NewBatchHandler[T any](l logger.Logger, fn func(mes []*sarama.ConsumerMessage, event []T) error) *BatchHandler[T] {
	return &BatchHandler[T]{l: l, fn: fn}
}

func (bh *BatchHandler[T]) Setup(session sarama.ConsumerGroupSession) error {
	return nil
}

func (bh *BatchHandler[T]) Cleanup(session sarama.ConsumerGroupSession) error {
	return nil
}

func (bh *BatchHandler[T]) ConsumeClaim(session sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) error {
	ch := claim.Messages()
	//批量处理10条消息
	batchSize := 10
	for {
		msgs := make([]*sarama.ConsumerMessage, 0, batchSize)
		ts := make([]T, 0, batchSize)
		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
		done := false
		for i := 0; i < batchSize && !done; i++ {
			select {
			case <-ctx.Done():
				// 这一批次已经超时了，
				// 或者，整个 consumer 被关闭了
				// 不再尝试凑够一批了
				done = true
			case msg, ok := <-ch:
				if !ok {
					cancel()
					// channel 被关闭了
					return nil
				}
				var v T
				err := json.Unmarshal(msg.Value, v)
				if err != nil {
					bh.l.Error("反序列化失败", logger.Field{
						Key: "Topic",
						Val: msg.Topic,
					}, logger.Field{
						Key: "Partition",
						Val: msg.Partition,
					}, logger.Field{
						Key: "offset",
						Val: msg.Offset,
					})

					continue
				}
				msgs = append(msgs, msg)
				ts = append(ts, v)
			}
		}
		cancel()
		//凑够一批了，可以处理
		err := bh.fn(msgs, ts)
		if err != nil {
			bh.l.Error("批量消费消息失败", logger.Field{
				Key: "error",
				Val: err,
			})
		}

		for _, msg := range msgs {
			session.MarkMessage(msg, "")
		}
	}
}
